from StringIO import StringIO
import numpy as np
import struct
import wave

def read_signal_from_wav_file(file_name):
    '''Reads the audio signal and it's framerate from a WAV file.
    
    Input:
        file_name - WAV file name
    Output:
        signal - audio signal as numpy array
        framerate - audio signal frame rate
    '''
    file = open(file_name, 'rb')
    wave_read = wave.open(file, 'r')
    return read_signal_from_waveread(wave_read)

def read_signal_from_string(string):
    '''Reads the audio signal and it's framerate from a WAV file as a string of bytes.
    
    Input:
        string - WAV file as a string of bytes
    Output:
        signal - audio signal as numpy array
        framerate - audio signal frame rate
    '''
    wave_read = wave.open(StringIO(string), 'r')
    return read_signal_from_waveread(wave_read)
    
def read_signal_from_waveread(wr):
    '''
    Reads the audio signal and it's framerate from the given Wave_Read object.
    
    Input:
        wr - Wave_read object
    Output:
        signal - audio signal as numpy array
        framerate - audio signal frame rate
    '''
    nframes = wr.getnframes()
    frames = wr.readframes(nframes)
    signal = str2nparray(frames, wr.getsampwidth(), nframes, wr.getnchannels())
    framerate = wr.getframerate()
    return signal, framerate

def str2nparray(str, samplewidth, nframes, nchannels):
    '''Converts the given string of bytes into a numpy array.
    
    str - the string to be converted
    samplewidth - sample width in bytes
    nframes - number of frames
    nchannels - number of channels
    '''
    types = {2:'h', 4:'i', 8:'l'}
    format = '<' + "{0}".format(nframes * nchannels) + types.get(samplewidth)
    samples = struct.unpack(format, str)
    return np.asarray(samples)

def extract_channel_data_from_signal(signal, channel, nchannels, frameWidth):
    sampleWidth = frameWidth / nchannels
    nframes = len(signal) / frameWidth
    data = []
    indexInFrame = channel * sampleWidth
    for i in range(nframes):
        frame = signal[i:i + frameWidth]
        sample = frame[indexInFrame:indexInFrame + sampleWidth]
        data.append(sample)
    return ''.join(data)

def extract_channel_data_from_wave_read(wave_read, channel):
    sampleWidth = wave_read.getsampwidth()
    indexInFrame = channel * sampleWidth
    data = []
    wave_read.setpos(0)
    frame = wave_read.readframes(1)
    while len(frame) > 0:
        sample = frame[indexInFrame:indexInFrame + sampleWidth]
        data.append(sample)
        frame = wave_read.readframes(1)
    return ''.join(data)